import torch
from kan import *
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import moviepy.video.io.ImageSequenceClip

if torch.cuda.is_available():
  device = torch.device("cuda")
else:
  device = torch.device("cpu")

# ========== 集成team-daniel/KAN回归 ==========
import sys
sys.path.append(r'../KAN')  # 假设KAN.py在上级目录的KAN文件夹下，如有需要请调整路径
from kan import *
import torch

# 数据准备
X = data[['freq_all', 'flow', 'dt']].values.astype(np.float32)
y = data['heat_load'].values.astype(np.float32).reshape(-1, 1)
X_tensor = torch.tensor(X)
y_tensor = torch.tensor(y)

# 构建KAN回归模型
model = KAN(width=[3, 10, 1])  # 输入3维，隐藏层10，输出1维
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = torch.nn.MSELoss()

# 训练
for epoch in range(1000):
    optimizer.zero_grad()
    y_pred = model(X_tensor)
    loss = loss_fn(y_pred, y_tensor)
    loss.backward()
    optimizer.step()
    if epoch % 100 == 0:
        print(f'KAN Epoch {epoch}, Loss: {loss.item()}')

# 预测与评价
with torch.no_grad():
    y_pred_kan = model(X_tensor).numpy().flatten()
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error, max_error
print('\n【team-daniel/KAN回归结果】')
print('KAN R2:', r2_score(y, y_pred_kan))
print('KAN MSE:', mean_squared_error(y, y_pred_kan))
print('KAN MAE:', mean_absolute_error(y, y_pred_kan))
print('KAN Max Error:', max_error(y, y_pred_kan))

# 导出KAN公式表达
print('\nKAN 公式表达:')
try:
    formula_str = model.get_formula(input_names=['freq_all', 'flow', 'dt'])
    print(formula_str)
    # 可选：导出为Python函数
    with open('kan_formula.py', 'w', encoding='utf-8') as f:
        f.write('def kan_predict(freq_all, flow, dt):\n')
        f.write(f'    return {formula_str}\n')
    print('公式已导出到 kan_formula.py')
except Exception as e:
    print('公式导出失败:', e)

